package com.agent.platform.service;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.*;
import java.util.stream.Collectors;

/**
 * Spring AI 模型选择器服务
 * 对应Node.js版本的OpenRouterModelSelector
 */
@Slf4j
@Service
public class ModelSelectorService {

    @Value("${spring.ai.openai.api-key}")
    private String apiKey;

    @Value("${spring.ai.openai.base-url:https://openrouter.ai/api/v1}")
    private String baseUrl;

    @Value("${agent.platform.model.default-model:anthropic/claude-3.5-sonnet}")
    private String defaultModel;

    @Value("${agent.platform.model.timeout:60000}")
    private int timeout;

    /**
     * 模型缓存配置 - 对应Node.js版本的modelCacheConfigs
     */
    private static final Map<String, CacheConfig> MODEL_CACHE_CONFIGS = Map.of(
        "anthropic/claude-sonnet-4", new CacheConfig(true, true, true, "ephemeral")
        // 可以添加更多模型配置
    );

    /**
     * 选择模型
     */
    public String selectModel(Map<String, Object> metadata) {
        if (metadata != null && metadata.containsKey("model")) {
            return (String) metadata.get("model");
        }
        return defaultModel;
    }

    /**
     * 调用模型 - 非流式
     */
    public ChatResponse callModel(
        List<Message> messages,
        List<Map<String, Object>> tools,
        String model
    ) {
        String modelToUse = model != null ? model : defaultModel;
        log.debug("Calling model: {}", modelToUse);

        try {
            // 创建OpenAI API客户端 (兼容OpenRouter)
            OpenAiApi openAiApi = new OpenAiApi(baseUrl, apiKey);

            // 配置选项
            OpenAiChatOptions options = OpenAiChatOptions.builder()
                .withModel(modelToUse)
                .withMaxTokens(4096)
                .withTemperature(0.7)
                .build();

            // 应用工具
            if (tools != null && !tools.isEmpty()) {
                // 转换工具格式
                options = OpenAiChatOptions.builder()
                    .withModel(modelToUse)
                    .withMaxTokens(4096)
                    .withTemperature(0.7)
                    .build();
            }

            // 创建ChatModel
            ChatModel chatModel = new OpenAiChatModel(openAiApi, options);

            // 应用缓存控制 (如果支持)
            List<Message> processedMessages = applyCacheControl(messages, modelToUse);

            // 调用模型
            Prompt prompt = new Prompt(processedMessages, options);
            ChatResponse response = chatModel.call(prompt);

            log.debug("Model response received. Tokens: {}",
                response.getMetadata().getUsage().getTotalTokens());

            return response;

        } catch (Exception e) {
            log.error("Model call failed for model: {}", modelToUse, e);
            throw new RuntimeException("Model call failed: " + e.getMessage(), e);
        }
    }

    /**
     * 流式调用模型
     */
    public Flux<ChatResponse> callModelStreaming(
        List<Message> messages,
        List<Map<String, Object>> tools,
        String model
    ) {
        String modelToUse = model != null ? model : defaultModel;
        log.debug("Calling model (streaming): {}", modelToUse);

        try {
            OpenAiApi openAiApi = new OpenAiApi(baseUrl, apiKey);

            OpenAiChatOptions options = OpenAiChatOptions.builder()
                .withModel(modelToUse)
                .withMaxTokens(4096)
                .withTemperature(0.7)
                .build();

            ChatModel chatModel = new OpenAiChatModel(openAiApi, options);

            List<Message> processedMessages = applyCacheControl(messages, modelToUse);
            Prompt prompt = new Prompt(processedMessages, options);

            return chatModel.stream(prompt);

        } catch (Exception e) {
            log.error("Streaming model call failed for model: {}", modelToUse, e);
            return Flux.error(new RuntimeException("Streaming call failed: " + e.getMessage(), e));
        }
    }

    /**
     * 应用缓存控制 - 对应Node.js版本的applyCacheControl
     */
    private List<Message> applyCacheControl(List<Message> messages, String model) {
        CacheConfig cacheConfig = MODEL_CACHE_CONFIGS.get(model);
        if (cacheConfig == null) {
            return messages;
        }

        log.debug("Applying cache control for model: {}", model);

        // Spring AI暂不直接支持缓存控制标记
        // 这里保留原始消息，实际缓存由OpenRouter/Anthropic后端处理
        // 如需自定义缓存逻辑，可在此扩展

        return messages;
    }

    /**
     * 健康检查
     */
    public boolean healthCheck() {
        try {
            // 简单的健康检查 - 尝试创建API客户端
            new OpenAiApi(baseUrl, apiKey);
            return true;
        } catch (Exception e) {
            log.error("Health check failed", e);
            return false;
        }
    }

    /**
     * 缓存配置内部类
     */
    private record CacheConfig(
        boolean systemMessage,
        boolean lastTool,
        boolean lastMessage,
        String cacheType
    ) {}
}
